package com.jshan.handler;

import com.jcraft.jsch.ChannelShell;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.Session;
import com.jshan.base.entity.SysServer;
import com.jshan.service.SysServerService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import cn.dev33.satoken.stp.StpUtil;
import cn.dev33.satoken.exception.SaTokenException;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.nio.file.Files;
import java.nio.file.Path;
import java.io.File;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

@Component
@Slf4j
public class TerminalWebSocketHandler extends TextWebSocketHandler {

    private final SysServerService serverService;

    // 存储 WebSocketSession 对应的 SSH Session 和 ChannelShell
    private final Map<String, Session> sshSessions = new ConcurrentHashMap<>();
    private final Map<String, ChannelShell> sshChannels = new ConcurrentHashMap<>();

    public TerminalWebSocketHandler(SysServerService serverService) {
        this.serverService = serverService;
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        log.info("WebSocket 连接建立: {}", session.getId());

        String serverIdStr = null;
        String token = null;

        try {
            String query = session.getUri().getQuery();
            if (query != null) {
                Pattern serverIdPattern = Pattern.compile("serverId=(\\d+)");
                Matcher serverIdMatcher = serverIdPattern.matcher(query);
                if (serverIdMatcher.find()) {
                    serverIdStr = serverIdMatcher.group(1);
                }

                Pattern tokenPattern = Pattern.compile("token=([^&]+)");
                Matcher tokenMatcher = tokenPattern.matcher(query);
                if (tokenMatcher.find()) {
                    token = tokenMatcher.group(1);
                }
            }
        } catch (Exception e) {
            log.error("无法从 WebSocket URI 获取参数: {}", e.getMessage());
            session.sendMessage(new TextMessage("Error: 无法获取连接参数。"));
            session.close(CloseStatus.SERVER_ERROR);
            return;
        }

        if (serverIdStr == null) {
            session.sendMessage(new TextMessage("Error: 未提供服务器ID。"));
            session.close(CloseStatus.SERVER_ERROR);
            return;
        }
        Long serverId = Long.parseLong(serverIdStr);

        // 1. Sa-Token 权限验证 (使用传递过来的 token)
        Long loginId = null;
        if (token != null && !token.isEmpty()) {
            try {
                Object loginIdObj = StpUtil.getLoginIdByToken(token);
                if (loginIdObj == null) {
                    throw new SaTokenException("Token 对应的登录ID为空，认证失败");
                }
                loginId = Long.parseLong(String.valueOf(loginIdObj));
                // Token 有效，现在检查权限
                if (!StpUtil.hasPermission(loginId, "server:terminal:connect")) {
                    log.warn("用户 {} 没有终端连接权限", loginId);
                    session.sendMessage(new TextMessage("Error: 您没有终端连接权限。"));
                    session.close(CloseStatus.SERVER_ERROR.withReason("Permission Denied"));
                    return;
                }
            } catch (SaTokenException e) {
                log.warn("Sa-Token 验证失败 (无效或过期令牌): {}", e.getMessage());
                session.sendMessage(new TextMessage("Error: 认证失败: 无效或过期令牌。"));
                session.close(CloseStatus.SERVER_ERROR.withReason("Authentication Failed"));
                return;
            } catch (Exception e) { // 捕获其他未知异常
                log.error("Sa-Token 验证过程中发生未知错误: {}", e.getMessage(), e);
                session.sendMessage(new TextMessage("Error: 认证过程中发生未知错误。"));
                session.close(CloseStatus.SERVER_ERROR.withReason("Authentication Error"));
                return;
            }
        } else {
            log.warn("未提供认证令牌，WebSocket 连接关闭");
            session.sendMessage(new TextMessage("Error: 未提供认证令牌。"));
            session.close(CloseStatus.SERVER_ERROR.withReason("Token Missing"));
            return;
        }

        SysServer server = serverService.getById(serverId);
        if (server == null) {
            session.sendMessage(new TextMessage("Error: 服务器不存在。"));
            session.close(CloseStatus.SERVER_ERROR);
            return;
        }

        try {
            JSch jsch = new JSch();
            Session sshSession = null;

            // 优先使用私钥
            if (server.getPrivateKey() != null && !server.getPrivateKey().trim().isEmpty()) {
                Path tempKeyPath = Files.createTempFile("ssh_key_", ".pem");
                Files.write(tempKeyPath, server.getPrivateKey().getBytes());
                File tempKeyFile = tempKeyPath.toFile();
                if (tempKeyFile.exists()) {
                    // 将文件权限设置为 600 (-rw-------)
                    // JSch 要求私钥文件权限严格，不能被其他用户读写
                    tempKeyFile.setReadable(false, false); // Remove read for others
                    tempKeyFile.setWritable(false, false); // Remove write for others
                    tempKeyFile.setExecutable(false, false); // Remove exec for others
                    tempKeyFile.setReadable(true, true); // Add read for owner
                    tempKeyFile.setWritable(true, true); // Add write for owner
                }

                jsch.addIdentity(tempKeyPath.toString());
                sshSession = jsch.getSession(server.getUsername(), server.getIp(), server.getPort());
                log.info("尝试使用私钥连接 SSH.");

            } else if (server.getPassword() != null && !server.getPassword().trim().isEmpty()) {
                sshSession = jsch.getSession(server.getUsername(), server.getIp(), server.getPort());
                sshSession.setPassword(server.getPassword());
                log.info("尝试使用密码连接 SSH.");
            } else {
                session.sendMessage(new TextMessage("Error: 未提供有效的SSH凭证 (密码或私钥)。"));
                session.close(CloseStatus.SERVER_ERROR);
                return;
            }

            sshSession.setConfig("StrictHostKeyChecking", "no");
            sshSession.connect(10000); // 连接超时 10 秒

            ChannelShell channel = (ChannelShell) sshSession.openChannel("shell");
            channel.connect(5000); // 通道连接超时 5 秒

            sshSessions.put(session.getId(), sshSession);
            sshChannels.put(session.getId(), channel);

            // 启动一个线程读取SSH输出并发送给WebSocket客户端
            new Thread(() -> {
                try {
                    InputStream in = channel.getInputStream();
                    byte[] buffer = new byte[1024];
                    int n;
                    while ((n = in.read(buffer)) > 0) {
                        session.sendMessage(new TextMessage(new String(buffer, 0, n)));
                    }
                } catch (IOException e) {
                    // 当 WebSocket 连接关闭时，可能会抛出 IOException
                    if (session.isOpen()) { // 只有当 WebSocket 仍在打开状态才记录错误
                        log.error("读取SSH输出失败或 WebSocket 已关闭: {}", e.getMessage());
                    }
                } finally {
                    // 当SSH连接关闭时，也关闭WebSocket连接
                    try {
                        if (session.isOpen()) {
                            session.close();
                        }
                    } catch (IOException e) {
                        log.error("关闭WebSocketSession失败: {}", e.getMessage());
                    }
                }
            }).start();

            // 发送欢迎信息或首次提示符
            session.sendMessage(new TextMessage("Connected to " + server.getName() + "\n"));

        } catch (Exception e) {
            log.error("SSH 连接失败: {}", e.getMessage());
            session.sendMessage(new TextMessage("SSH connection failed: " + e.getMessage()));
            session.close(CloseStatus.SERVER_ERROR);
        }
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        ChannelShell channel = sshChannels.get(session.getId());
        if (channel != null && channel.isConnected()) { // 确保通道仍然连接
            try {
                OutputStream out = channel.getOutputStream();
                out.write(message.getPayload().getBytes());
                out.flush();
            } catch (IOException e) {
                log.error("写入SSH命令失败: {}", e.getMessage());
                // 可以在这里通知前端命令执行失败
                session.sendMessage(new TextMessage("\r\nError: 无法发送命令到SSH。\r\n"));
            }
        } else {
            log.warn("尝试向已关闭或不存在的 SSH 通道发送命令: {}", session.getId());
            session.sendMessage(new TextMessage("\r\nError: 终端未连接或已关闭，无法发送命令。\r\n"));
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        log.info("WebSocket 连接关闭: {} - {}, Reason: {}", session.getId(), status.getCode(), status.getReason());

        ChannelShell channel = sshChannels.remove(session.getId());
        if (channel != null) {
            channel.disconnect();
            log.info("SSH Channel 断开连接: {}", session.getId());
        }

        Session sshSession = sshSessions.remove(session.getId());
        if (sshSession != null) {
            sshSession.disconnect();
            log.info("SSH Session 断开连接: {}", session.getId());
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        log.error("WebSocket 传输错误: {}, Session ID: {}", exception.getMessage(), session.getId());
        // 尝试关闭连接，如果连接仍然打开
        if (session.isOpen()) {
            session.close(CloseStatus.SERVER_ERROR.withReason("传输错误: " + exception.getMessage()));
        }
        afterConnectionClosed(session, CloseStatus.SERVER_ERROR); // 确保资源被清理
    }
}
